#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
RAG智能对话机器人演示
展示基于Milvus和Ollama的完整RAG功能
"""

import time
from tools.RAGChatBot import RAGChatBot


def print_separator(title="", char="=", length=80):
    """打印分隔线"""
    if title:
        title_len = len(title)
        side_len = (length - title_len - 2) // 2
        print(char * side_len + f" {title} " + char * side_len)
    else:
        print(char * length)


def demo_basic_chat():
    """基础对话演示"""
    print_separator("基础对话演示")
    
    # 初始化RAG机器人
    print("🚀 初始化RAG系统...")
    bot = RAGChatBot(
        # ollama_url="http://172.26.32.1:11434",
        ollama_url="http://127.0.0.1:11434",
        milvus_collection="massive_file_search"
    )
    
    if not bot.initialize():
        print("❌ 初始化失败")
        return None
    
    # 演示问题列表
    demo_questions = [
        "有哪些Python文件？",
        "最近修改的配置文件是什么？",
        "帮我找找文档类的文件",
        "有关于数据库的文件吗？",
        "项目中有README文件吗？"
    ]
    
    print("\n💬 开始演示对话:")
    print("-" * 60)
    
    for i, question in enumerate(demo_questions, 1):
        print(f"\n[问题 {i}] 👤: {question}")
        print("🔍 正在搜索和生成回答...")
        
        start_time = time.time()
        response, sources = bot.chat(question)
        elapsed = time.time() - start_time
        
        print(f"🤖: {response}")
        
        if sources:
            print(f"\n📚 参考文件 ({len(sources)}个):")
            for j, source in enumerate(sources[:3], 1):  # 只显示前3个
                print(f"   {j}. {source['file_name']} (相似度: {source['similarity_score']:.3f})")
        
        print(f"⏱️ 耗时: {elapsed:.2f}秒")
        print("-" * 60)
        
        # 短暂停顿
        time.sleep(1)
    
    return bot


def demo_model_switching(bot):
    """模型切换演示"""
    if not bot:
        return
    
    print_separator("模型切换演示")
    
    # 获取可用模型
    models = bot.get_available_models()
    print(f"📋 可用模型 ({len(models)}个):")
    for i, model in enumerate(models, 1):
        current = " (当前)" if model == bot.current_model else ""
        print(f"   {i}. {model}{current}")
    
    if len(models) > 1:
        # 尝试切换到另一个模型
        other_models = [m for m in models if m != bot.current_model]
        if other_models:
            new_model = other_models[0]
            print(f"\n🔄 尝试切换到: {new_model}")
            
            if bot.switch_model(new_model):
                print("✅ 切换成功")
                
                # 用新模型回答一个问题
                question = "用新模型回答：项目中有什么类型的文件？"
                print(f"\n👤: {question}")
                
                response, sources = bot.chat(question)
                print(f"🤖: {response}")
            else:
                print("❌ 切换失败")


def demo_configuration(bot):
    """配置演示"""
    if not bot:
        return
    
    print_separator("配置管理演示")
    
    # 显示当前配置
    stats = bot.get_stats()
    print("⚙️ 当前配置:")
    print(f"   当前模型: {stats['current_model']}")
    print(f"   最大检索文件数: {stats['config']['max_context_files']}")
    print(f"   最大上下文长度: {stats['config']['max_context_length']}")
    print(f"   相似度阈值: {stats['config']['similarity_threshold']}")
    
    # 修改配置
    print("\n🔧 调整配置参数...")
    bot.update_config(
        max_context_files=3,
        similarity_threshold=0.4
    )
    
    # 显示新配置
    new_stats = bot.get_stats()
    print("\n⚙️ 新配置:")
    print(f"   最大检索文件数: {new_stats['config']['max_context_files']}")
    print(f"   相似度阈值: {new_stats['config']['similarity_threshold']}")


def demo_file_search(bot):
    """文件搜索演示"""
    if not bot:
        return
    
    print_separator("文件搜索演示")
    
    search_queries = [
        "python",
        "config", 
        "document",
        "test",
        "data"
    ]
    
    for query in search_queries:
        print(f"\n🔍 搜索关键词: '{query}'")
        
        results = bot.search_files(query, top_k=5)
        
        if results:
            print(f"📄 找到 {len(results)} 个相关文件:")
            for i, result in enumerate(results, 1):
                print(f"   {i}. {result['file_name']}")
                print(f"      路径: {result['file_path']}")
                print(f"      相似度: {result['similarity_score']:.3f}")
                print(f"      类型: {result['file_type']}")
                print()
        else:
            print("❌ 未找到相关文件")


def demo_chat_history(bot):
    """对话历史演示"""
    if not bot:
        return
    
    print_separator("对话历史管理演示")
    
    # 显示对话历史
    history = bot.get_chat_history()
    print(f"📝 对话历史 ({len(history)}条消息):")
    
    if history:
        for i, msg in enumerate(history[-6:], 1):  # 显示最后6条
            role = "👤" if msg.role == "user" else "🤖"
            content = msg.content[:100] + "..." if len(msg.content) > 100 else msg.content
            timestamp = time.strftime("%H:%M:%S", time.localtime(msg.timestamp))
            print(f"   [{timestamp}] {role}: {content}")
    else:
        print("   暂无对话记录")
    
    # 清空历史
    print("\n🗑️ 清空对话历史...")
    bot.clear_history()
    
    # 确认清空
    new_history = bot.get_chat_history()
    print(f"✅ 历史已清空，当前消息数: {len(new_history)}")


def demo_performance_analysis(bot):
    """性能分析演示"""
    if not bot:
        return
    
    print_separator("性能分析演示")
    
    test_questions = [
        "找找所有的Python文件",
        "有什么配置相关的文件？",
        "最近的文档更新了什么？"
    ]
    
    performance_data = []
    
    print("📊 性能测试进行中...")
    
    for i, question in enumerate(test_questions, 1):
        print(f"\n[测试 {i}] {question}")
        
        # 记录各阶段时间
        total_start = time.time()
        
        # 搜索阶段
        search_start = time.time()
        retrieved_files = bot.search_files(question)
        search_time = time.time() - search_start
        
        # 生成回答阶段
        if retrieved_files:
            context = bot.build_context(question, retrieved_files)
            
            generate_start = time.time()
            response = bot.generate_response(question, context)
            generate_time = time.time() - generate_start
        else:
            generate_time = 0
            response = "未找到相关文件"
        
        total_time = time.time() - total_start
        
        # 记录性能数据
        perf_data = {
            'question': question,
            'search_time': search_time,
            'generate_time': generate_time,
            'total_time': total_time,
            'files_found': len(retrieved_files),
            'response_length': len(response)
        }
        performance_data.append(perf_data)
        
        # 显示结果
        print(f"   🔍 搜索耗时: {search_time:.3f}秒")
        print(f"   🤖 生成耗时: {generate_time:.3f}秒")
        print(f"   ⏱️ 总耗时: {total_time:.3f}秒")
        print(f"   📄 找到文件: {len(retrieved_files)}个")
    
    # 性能总结
    print("\n📈 性能总结:")
    avg_search = sum(p['search_time'] for p in performance_data) / len(performance_data)
    avg_generate = sum(p['generate_time'] for p in performance_data) / len(performance_data)
    avg_total = sum(p['total_time'] for p in performance_data) / len(performance_data)
    
    print(f"   平均搜索时间: {avg_search:.3f}秒")
    print(f"   平均生成时间: {avg_generate:.3f}秒") 
    print(f"   平均总时间: {avg_total:.3f}秒")


def main():
    """主演示函数"""
    print_separator("RAG智能对话机器人演示", "=", 80)
    print("🎯 演示基于Milvus向量检索和Ollama大模型的RAG系统")
    print("📋 演示内容:")
    print("   1. 基础对话功能")
    print("   2. 模型切换")
    print("   3. 配置管理")
    print("   4. 文件搜索")
    print("   5. 对话历史管理")
    print("   6. 性能分析")
    print()
    
    try:
        # 运行基础对话演示
        bot = demo_basic_chat()
        
        if bot:
            # 运行其他演示
            demo_model_switching(bot)
            demo_configuration(bot)
            demo_file_search(bot)
            demo_chat_history(bot)
            demo_performance_analysis(bot)
            
            # 最终统计
            print_separator("最终统计")
            final_stats = bot.get_stats()
            print(f"🏆 演示完成:")
            print(f"   使用模型: {final_stats['current_model']}")
            print(f"   处理消息: {final_stats['chat_messages']}条")
            print(f"   文件索引: {final_stats['milvus_entities']:,}条")
        
        print_separator("演示结束", "=", 80)
        
    except KeyboardInterrupt:
        print("\n⏹️ 演示被中断")
    except Exception as e:
        print(f"❌ 演示过程中出错: {e}")


if __name__ == "__main__":
    main()